import torch
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
from pystct import isdct



def SNR(cover, container, phase, transform, transform_constructor=None, on_phase=False):
	"""
    Computes SNR (Signal-to-Noise-Ratio)
	metric between cover and container signals.
	First, it computes i[transform] over the spectrograms.
	Transform can be either [cosine] or [fourier]

	" A local SNR of 30dB is effectively a clean signal. 
	Listeners will barely notice anything better than 20dB, 
	and intelligibility is still pretty good at 0dB SNR "
	> http://www1.icsi.berkeley.edu/Speech/faq/speechSNR.html
    """
	if transform == 'cosine':
		cover_wav = isdct(cover.squeeze(0).squeeze(0).cpu().detach().numpy(), frame_step=62)
		noise_wav = isdct((container - cover).squeeze(0).squeeze(0).cpu().detach().numpy(), frame_step=62)
	elif (transform == 'fourier') and (transform_constructor is not None):
		if on_phase:
			cover_wav = transform_constructor.inverse(cover.squeeze(1), phase.squeeze(1)).cpu().data.numpy()[..., :]
			noise_wav = transform_constructor.inverse(cover.squeeze(1), (container - phase).squeeze(1)).cpu().data.numpy()[..., :]
		else:
			cover_wav = transform_constructor.inverse(cover.squeeze(1), phase.squeeze(1)).cpu().data.numpy()[..., :]
			noise_wav = transform_constructor.inverse((container - cover).squeeze(1), phase.squeeze(1)).cpu().data.numpy()[..., :]
	else: raise Exception('Transform not defined')
	
	signal = np.sum(np.abs(np.fft.fft(cover_wav)) ** 2) / len(np.fft.fft(cover_wav))
	noise = np.sum(np.abs(np.fft.fft(noise_wav))**2) / len(np.fft.fft(noise_wav))
	if noise <= 0.00001 or signal <= 0.00001: return 0
	return 10 * np.log10(signal / noise)

# 修改后的StegoLoss函数，计算宿主音频与容器音频、隐写音频与恢复音频之间的损失
def StegoLoss(cover : torch.Tensor, container : torch.Tensor, secret : torch.Tensor, revealed : torch.Tensor, beta : float = 0.5) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Custom StegoLoss function: a convex combination of two losses:
    - Loss between cover and container audios (L_cover)
    - Loss between secret and revealed audios (L_secret)
    Both terms are leveraged with the hyperparameter beta.
    """
    # 宿主与容器音频之间的损失（均方误差）
    loss_cover = F.mse_loss(cover, container)
    
    # 隐写与恢复音频之间的损失（均方误差）
    loss_secret = F.mse_loss(secret, revealed)
    
    # 结合两部分的损失
    loss = (1 - beta) * loss_cover + beta * loss_secret
    
    return loss, loss_cover, loss_secret


# 评价指标计算，保留与音频相关的评价指标
def calculate_metrics(cover : torch.Tensor, container : torch.Tensor, secret : torch.Tensor, revealed : torch.Tensor) -> dict[str, torch.Tensor]:
    """
    Function to calculate metrics for evaluation
    """
    # 计算SNR
    snr_value = SNR(cover, container)
    
    # 计算宿主与容器音频之间的损失
    loss_cover = F.mse_loss(cover, container)
    
    # 计算隐写与恢复音频之间的损失
    loss_secret = F.mse_loss(secret, revealed)
    
    return {
        'SNR': snr_value,
        'Loss_Cover': loss_cover,
        'Loss_Secret': loss_secret
    }